package top.zhangjianyong.tools.utils;

import lombok.extern.slf4j.Slf4j;


// 线性回归模型类，用于拟合数据并计算回归系数和统计指标
@Slf4j
public class LinearRegression {

    // 成员变量添加注释说明
    private double intercept; // 线性回归方程的截距项
    private double slope;     // 线性回归方程的斜率系数
    private double r2;        // 决定系数R²（拟合优度指标）
    private double svar0;     // 截距项的标准误差方差
    private double svar1;     // 斜率项的标准误差方差

    // 构造函数：根据输入数据x和y初始化回归模型
    public LinearRegression(double[] x, double[] y) {
        if (x.length != y.length) {
            throw new IllegalArgumentException("数组长度不一致");
        }
        int n = x.length;

        // 第一次遍历计算x和y的总和及x的平方和
        double sumx = 0.0, sumy = 0.0;
        for (int i = 0; i < n; i++) {
            sumx += x[i];
            sumy += y[i];
        }
        double xbar = sumx / n;    // x的均值
        double ybar = sumy / n;    // y的均值

        // 第二次遍历计算协方差和方差
        double xxbar = 0.0, yybar = 0.0, xybar = 0.0;
        for (int i = 0; i < n; i++) {
            xxbar += (x[i] - xbar) * (x[i] - xbar);  // x的方差分母
            yybar += (y[i] - ybar) * (y[i] - ybar);  // y的总平方和（TSS）
            xybar += (x[i] - xbar) * (y[i] - ybar);  // x与y的协方差
        }
        slope = xybar / xxbar;     // 斜率计算公式：协方差/方差
        intercept = ybar - slope * xbar; // 截距公式：y均值 - 斜率*x均值

        // 计算残差平方和（RSS）和回归平方和（SSR）
        double rss = 0.0;
        double ssr = 0.0;
        for (int i = 0; i < n; i++) {
            double fit = slope * x[i] + intercept;
            rss += (fit - y[i]) * (fit - y[i]);  // 残差平方和
            ssr += (fit - ybar) * (fit - ybar);   // 回归平方和
        }

        int degreesOfFreedom = n - 2; // 自由度计算（n-2）
        r2 = ssr / yybar;             // R²决定系数 = SSR/TSS
        double svar = rss / degreesOfFreedom; // 残差方差估计
        svar1 = svar / xxbar;         // 斜率标准误差的方差
        svar0 = svar / n + xbar * xbar * svar1; // 截距标准误差的方差
    }

    // 获取回归方程的截距项
    public double intercept() {
        return intercept;
    }

    // 获取回归方程的斜率
    public double slope() {
        return slope;
    }

    // 获取决定系数R²（拟合优度）
    public double R2() {
        return r2;
    }

    // 获取截距的标准误差
    public double interceptStdErr() {
        return Math.sqrt(svar0);
    }

    // 获取斜率的标准误差
    public double slopeStdErr() {
        return Math.sqrt(svar1);
    }

    // 根据输入x值预测y值
    public double predict(double x) {
        return slope * x + intercept;
    }

    // 主函数：测试用例
    public static void main(String[] args) {
        double[] x = {1, 2, 3, 4, 5};
        double[] y = {110, 100, 129, 80, 130};

        LinearRegression lr = new LinearRegression(x, y);
        System.out.println("y   = " + lr.slope() + " * x + " + lr.intercept());
        System.out.println("R^2 = " + lr.R2());
        System.out.println("Slope standard error = " + lr.slopeStdErr());
        System.out.println("Intercept standard error = " + lr.interceptStdErr());
        System.out.println("Predict y = " + lr.predict(8));
        System.out.println(" " + (lr.slope / lr.slopeStdErr()));
    }
}